package com.ai.module.mj.wss.handle;

import cn.hutool.core.text.CharSequenceUtil;
import com.ai.module.chat.common.constant.ChatImagePointStringConstant;
import com.ai.module.chat.common.constant.FileStringConstant;
import com.ai.module.chat.common.constant.StringConstant;
import com.ai.module.chat.common.domain.bo.AliUploadBo;
import com.ai.module.chat.common.domain.vo.AliFileVo;
import com.ai.module.business.service.BusinessTotalService;
import com.ai.module.business.service.ImageGroundService;
import com.ai.module.business.service.MemberCenterService;
import com.ai.module.mj.constant.enums.TaskStatus;
import com.ai.module.mj.constant.stringConstant.Constants;
import com.ai.module.mj.constant.stringConstant.MessageType;
import com.ai.module.mj.domain.Task;
import com.ai.module.mj.domain.bo.UserPointsBo;
import com.ai.module.mj.domain.dto.TaskCondition;
import com.ai.module.mj.loadbalancer.DiscordInstance;
import com.ai.module.mj.loadbalancer.DiscordLoadBalancer;
import com.ai.module.mj.service.HighTaskStoreService;
import com.ai.module.user.domain.entity.*;
import com.ai.module.user.enu.ImageType;
import com.ai.module.user.mapper.AIProductRecordMapper;
import com.ai.module.user.mapper.UserImageMapper;
import com.ai.module.user.service.UserBlendImageService;
import com.ai.module.user.service.UserImageAnalyzeService;
import com.ai.module.user.service.UserTaskService;
import com.ai.module.chat.util.DateUtils;
import com.ai.module.chat.util.SpringUtils;
import com.ai.module.chat.util.aliyun.AliyunUtil;
import com.ai.module.mj.constant.enums.TaskAction;
import com.ai.module.mj.domain.enu.ImageTaskSource;
import com.ai.module.mj.service.TranslateService;
import com.ai.module.mj.util.DiscordURLHelper;
import com.ai.module.websocket.service.AIProductWebSocketService;
import com.ai.module.websocket.service.GlobalWebSocketService;
import com.ai.module.websocket.vo.WebSocketProduct;
import com.alibaba.fastjson2.JSONObject;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.StringPool;
import lombok.extern.slf4j.Slf4j;
import net.dv8tion.jda.api.utils.data.DataArray;
import net.dv8tion.jda.api.utils.data.DataObject;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Async;
import org.springframework.util.CollectionUtils;
import org.springframework.web.socket.WebSocketSession;

import javax.annotation.Resource;
import java.time.LocalDateTime;
import java.util.*;

import static com.ai.module.chat.common.constant.FileStringConstant.MID_JOURNEY_AI;

@Slf4j
public abstract class MessageHandler {
	@Resource
	protected DiscordLoadBalancer discordLoadBalancer;

	public static HashMap<String, Task> storeTaskMap= new HashMap<>();
	@Resource
	protected DiscordURLHelper discordURLHelper;

	@Autowired
	private AIProductWebSocketService aiProductWebSocketService;

	private MemberCenterService memberCenterService = SpringUtils.getClazz(MemberCenterService.class);

	public abstract void handle(MessageType messageType, DataObject message);

	protected String getMessageContent(DataObject message) {
		return message.hasKey("content") ? message.getString("content") : "";
	}

	protected String getMessageNonce(DataObject message) {
		return message.hasKey("nonce") ? message.getString("nonce") : "";
	}

	protected void findAndFinishImageTask(TaskCondition condition, String finalPrompt, DataObject message) {
		String imageUrl = getImageUrl(message);
		String messageHash = this.discordURLHelper.getMessageHash(imageUrl);
		condition.setMessageHash(messageHash);
		Task task = this.discordLoadBalancer.findRunningTask(condition)
				.findFirst().orElseGet(() -> {
					condition.setMessageHash(null);
					return this.discordLoadBalancer.findRunningTask(condition)
							.min(Comparator.comparing(Task::getStartTime))
							.orElse(null);
				});
		if (task == null) {
//			log.error("任务已经不存在 [{}]",condition);
			return;
		}
		task.setProperty(Constants.TASK_PROPERTY_FINAL_PROMPT, finalPrompt);
		task.setProperty(Constants.TASK_PROPERTY_MESSAGE_HASH, messageHash);
		task.setImageUrl(imageUrl);
		task.setFinalPrompt(finalPrompt);
		// 释放锁
		finishTask(task, message);
		task.awake();

	}

	protected void finishTask(Task task, DataObject message) {
		// 进行 图片上传
		task.setProperty(Constants.TASK_PROPERTY_MESSAGE_ID, message.getString("id"));
		task.setProperty(Constants.TASK_PROPERTY_FLAGS, message.getInt("flags", 0));
		task.setProperty(Constants.TASK_PROPERTY_MESSAGE_HASH, this.discordURLHelper.getMessageHash(task.getImageUrl()));
//		task.success();
		Task successTask = new Task();
		BeanUtils.copyProperties(task,successTask);
		storeTaskMap.put(task.getId(),successTask);
		asyncUploadSave(task);
	}
	private HighTaskStoreService highTaskStoreService = SpringUtils.getClazz(HighTaskStoreService.class);
	// @Async // TODO
	public void asyncUploadSave(Task task) {
		try {
			ImageTaskSource taskSource = task.getTaskSource();
			if (taskSource==ImageTaskSource.commonImage ){
				 dealCommonImageTask(task);
			}else {
				AliFileVo aliFileVo=new AliFileVo();
				String fileName =task.getAction()+"_"+ DateUtils.getTransDateTime();
				// 处理 AI 应用 生成记录问题, 通用表
				AIProductRecordMapper aiProductRecord = SpringUtils.getClazz(AIProductRecordMapper.class);
				AIProductRecordEntity productRecord = new AIProductRecordEntity();
				productRecord.setBusinessId(task.getRobotId());
				productRecord.setUserId(task.getUserId());
				productRecord.setCostTime(System.currentTimeMillis()- task.getStart());
				TaskAction action = task.getAction();
				String realFileName = String.format("%s/%s/%s", FileStringConstant.AI_PRODUCT
						,MID_JOURNEY_AI,
						task.getUserIntId());
				String midJourneyUrl = task.getImageUrl();
				switch (action){
					case DESCRIBE:
						dealDesc(task);
						List<String> imageDescList = task.getImageDescList();
						productRecord.setAnswer(JSONObject.toJSONString(imageDescList));
						productRecord.setUserFile(
								JSONObject.toJSONString(Collections.singletonList(task.getImageUrl()))
						);
						break;
					case BLEND:
						productRecord.setQuestion(null);
//						productRecord.setUserFile();
					case IMAGINE:
					case UPSCALE:
						productRecord.setUserFile(
								JSONObject.toJSONString(task.getFileUrlList())
						);
						productRecord.setQuestion(task.getPrompt());
						 aliFileVo = AliyunUtil.uploadImageAndGetUrl(midJourneyUrl, String.format("%s/%s",realFileName,
								action+"_"+	DateUtils.getTransDateTime()));
						task.setScaleUrl(aliFileVo.getScaleFilePath());
						task.setImageUrl(aliFileVo.getFilePath());
//						task.setFileUrlList(Collections.singletonList(aliFileVo.getFilePath()));
						task.success();
						productRecord.setAiFile(JSONObject.toJSONString(Collections.singletonList(aliFileVo.getFilePath())));

				}
				productRecord.setAction(action);
				aiProductRecord.insert(productRecord);
				UserTaskEntity userTask = new UserTaskEntity();
				userTask.setStatus(TaskStatus.SUCCESS);
				userTask.setTaskId(task.getId());
				if (task.getAction() == TaskAction.IMAGINE
						&& !CollectionUtils.isEmpty(task.getFileUrlList()) ){
					String imageGroundPrompt = task.getImageGroundPrompt();
					userTask.setPrompt(imageGroundPrompt);
				}
				userTask.setStatus(task.getStatus());
				asyncInsert(task, userTask, fileName, aliFileVo);
			}
			dealImaginePoints(task);

		}catch (Exception ex){
			ex.printStackTrace();
			log.error("处理midJourney图片绘画失败[{}]", ex.getMessage());
		}finally {
			task.success();
			sendWebSocketImage(task);
			this.highTaskStoreService.save(task);

		}
	}

	private void dealCommonImageTask(Task task) {
		long start = System.currentTimeMillis();
		UserTaskEntity userTask = new UserTaskEntity();
		String fileName =task.getAction()+"_"+ DateUtils.getTransDateTime();
		// String.format("%s/%s", fileType, this.id)
		String realFileName = String.format("%s/%s/%s", FileStringConstant.MID_JOURNEY_IMAGE,MID_JOURNEY_AI, task.getUserIntId());
		String midJourneyUrl = task.getImageUrl();
		BeanUtils.copyProperties(task,userTask);
		TaskAction action = task.getAction();

		// 如果是 图片描述 就不需要 上传了
		if (action == TaskAction.DESCRIBE){
			dealImaginePoints(task);
			String descUrl = task.getDescription()
					.replaceAll("/describe", "")
					.replaceAll(" ", "");
			task.setFileUrlList(Collections.singletonList(descUrl));
			task.setImageUrl(descUrl);
			dealDesc(task);

		}else {
			AliFileVo aliFileVo = AliyunUtil.uploadImageAndGetUrl(midJourneyUrl, String.format("%s/%s",realFileName,fileName));
			task.setScaleUrl(aliFileVo.getScaleFilePath());
			task.setImageUrl(aliFileVo.getFilePath());
			task.success();
			try {
				task.notify(); // 为了任务退出
			}catch (IllegalMonitorStateException ex){
				ex.printStackTrace();
				log.error("唤醒有异常 IllegalMonitorStateException[{}]", ex.getMessage());
			}catch (Exception ex){
				ex.printStackTrace();
				log.error("唤醒有异常[{}]", ex.getMessage());
			}
			log.info("开始图片存储 耗时[{}]，Redis更新的数据为[{}], aliFileVo [{}]", System.currentTimeMillis()-start, task,aliFileVo);
			asyncInsert(task, userTask, fileName, aliFileVo);
		}
	}

	@Async
	public void dealImaginePoints(Task task) {
		UserPointsBo userPointsBo = new UserPointsBo();
		userPointsBo.setPrompt(task.getPrompt());
		userPointsBo.setUserId(task.getUserId());
		userPointsBo.setPointType(task.getAction());
		userPointsBo.setBusinessId(task.getUuid());
		String speedType = task.getSpeedType();
		String loginToken = task.getLoginToken();
		int costPoints = 0;
		if (StringConstant.FAST.equalsIgnoreCase(speedType)) {
			if (TaskAction.IMAGINE==task.getAction() ){
				costPoints = ChatImagePointStringConstant.fastCost;
			}else {
				costPoints = ChatImagePointStringConstant.fastCostUV;
			}
			// 快速模型
//			checkUserPoints(userLoginInfo, costPoints);

		} else {
			if (TaskAction.IMAGINE==task.getAction() ){
				costPoints = ChatImagePointStringConstant.slowCost;
			}else {
				costPoints = ChatImagePointStringConstant.slowCostUV;
			}
//			checkUserPoints(userLoginInfo, costPoints);
		}
		userPointsBo.setCost(costPoints);
		// 画前修改积分 todo
		memberCenterService.dealMidJourneyUserPoints(userPointsBo,task);
	}
	protected boolean hasImage(DataObject message) {
		DataArray attachments = message.optArray("attachments").orElse(DataArray.empty());
		return !attachments.isEmpty();
	}

	protected String getImageUrl(DataObject message) {
		DataArray attachments = message.getArray("attachments");
		if (!attachments.isEmpty()) {
			String imageUrl = attachments.getObject(0).getString("url");
			return replaceCdnUrl(imageUrl);
		}
		return null;
	}

	protected String replaceCdnUrl(String imageUrl) {
		if (CharSequenceUtil.isBlank(imageUrl)) {
			return imageUrl;
		}
		String cdn = this.discordURLHelper.getCdn();
		if (CharSequenceUtil.startWith(imageUrl, cdn)) {
			return imageUrl;
		}
		return CharSequenceUtil.replaceFirst(imageUrl, DiscordURLHelper.DISCORD_CDN_URL, cdn);
	}
	@Autowired
	private UserImageAnalyzeService imageAnalyze;

	 @Async
	public void asyncInsert(Task task,
							UserTaskEntity userTask,
							String fileName,
							AliFileVo aliFileVo) {
		// 处理 图片积分问题

		UserTaskService taskService =  SpringUtils.getClazz(UserTaskService.class);

		taskService
				.update(userTask,
				new LambdaQueryWrapper<UserTaskEntity>().eq(UserTaskEntity::getTaskId, task.getId()));
		UserImageEntity userImage = new UserImageEntity();
		BeanUtils.copyProperties(task,userImage);
		userImage.setFileName(fileName);
		userImage.setFilePath(aliFileVo.getFilePath());
		userImage.setExpireTime(aliFileVo.getExpireTime());
		userImage.setBusinessId(task.getId());
		userImage.setBucketName(FileStringConstant.USER_IMAGE);
		userImage.setScalePath(aliFileVo.getScaleFilePath());
		userImage.setProjectName(aliFileVo.getProjectName());
		SpringUtils.getClazz(UserImageMapper.class) .insert(userImage);
		//
		if (task.getAction() == TaskAction.UPSCALE) {
			ImageGroundEntity imageGround = new ImageGroundEntity();
			BeanUtils.copyProperties(userImage, imageGround);
			imageGround.setPrompt(task.getImageGroundPrompt());
			imageGround.setSourceType("CHINA");
			SpringUtils.getClazz(ImageGroundService.class) .save(imageGround);
			SpringUtils.getClazz(BusinessTotalService.class).saveTotalBusiness(imageGround.getUuid(), "IMAGE_GROUND");
		}
		// 进行 叠图数据转换 入库
		List<String> aliFileList = task.getFileUrlList();
		if (! CollectionUtils.isEmpty(aliFileList)){
			// BLEND/userId/
			UserBlendImageEntity userBlendImage = new UserBlendImageEntity();
			AliUploadBo aliUploadBo = new AliUploadBo();
			userBlendImage.setExpireTime(aliUploadBo.getExpireDate());
			userBlendImage.setUserId(task.getUserId());
			userBlendImage.setTaskId(task.getId());
			String imageUrlListStr = String.join(StringPool.COMMA,aliFileList);
			userBlendImage.setImageUrlList(imageUrlListStr);
			//
			UserBlendImageService imageService = SpringUtils.getClazz(UserBlendImageService.class);
			imageService.save(userBlendImage);
		}

	}

	private void dealDesc(Task task  ) {
		LocalDateTime expireTime = LocalDateTime.now().plusMonths(1);
		if (task.getAction()== TaskAction.DESCRIBE) {
			String prompt = task.getPrompt();
			TranslateService translateList = SpringUtils.getClazz(TranslateService.class);
			String promptZH = translateList.translateToChinese(prompt);

			String[] aiContentList = promptZH.replaceAll(".️⃣ ", ",").split("\n");
			String robotContent = String.join(",", aiContentList);
			task.setPrompt(promptZH);
			task.setImageDescList(Arrays.asList(aiContentList));
			UserImageAnalyzeEntity imageAnalyze1 = new UserImageAnalyzeEntity();
			imageAnalyze1.setRobotContent(robotContent);
			imageAnalyze1.setRobotId(task.getRobotId());
			imageAnalyze1.setFilePath(task.getImageUrl());
			imageAnalyze1.setSource(ImageType.description);
			imageAnalyze1.setPrompt(task.getImageUrl());
			imageAnalyze1.setUuid(task.getId());
			imageAnalyze1.setExpireTime(expireTime);
			imageAnalyze1.setUserId(task.getUserId());
			imageAnalyze.save(imageAnalyze1);

		}

	}
	public void sendWebSocketImage(Task task) {
		String userId = task.getUserId();
		String robotId = task.getRobotId();
		String userRobotKey = String.format(StringConstant.userROBOT, userId, robotId);
		Map<String, WebSocketSession> userProductSessionMap = GlobalWebSocketService.userProductSessionMap;
		if (userProductSessionMap.containsKey(userRobotKey)) {
			WebSocketSession webSocketSession = userProductSessionMap.get(userRobotKey);
			WebSocketProduct product = new WebSocketProduct();
			WebSocketProduct.AiImage aiImage = new WebSocketProduct.AiImage();
			aiImage.setProcess(task.getProgress());
			aiImage.setTask(task);
			aiImage.setTaskId(task.getId());
			product.setAiImage(aiImage);
			aiProductWebSocketService.sendMessage(webSocketSession,product);
		}
	}
	public void sendNoEndImage(Task task) {
		String userId = task.getUserId();
		String robotId = task.getRobotId();
		String userRobotKey = String.format(StringConstant.userROBOT, userId, robotId);
		Map<String, WebSocketSession> userProductSessionMap = GlobalWebSocketService.userProductSessionMap;
		if (userProductSessionMap.containsKey(userRobotKey)) {
			WebSocketSession webSocketSession = userProductSessionMap.get(userRobotKey);
			WebSocketProduct product = new WebSocketProduct();
			WebSocketProduct.AiImage aiImage = new WebSocketProduct.AiImage();
			aiImage.setProcess(task.getProgress());
//			aiImage.setTask(task);
			aiImage.setTaskId(task.getId());
			product.setAiImage(aiImage);
			aiProductWebSocketService.sendMessage(webSocketSession,product);
		}
	}
}
